import click
import mlflow
from ultralytics import YOLO
from predict_model import YOLOModel

@click.command()
@click.option("--training_data", "-d", type=str, default='', help="训练数据集路径")
@click.option("--imgsz", "-i", type=int, default=640, help="图片尺寸")
@click.option("--batch", "-b", type=int, default=16, help="批次大小")
@click.option("--epochs", "-e", type=int, default=50, help="训练轮数")
def train(training_data, imgsz, batch, epochs):
    # 记录模型参数
    mlflow.log_param("data", training_data)
    mlflow.log_param("imgsz", imgsz)
    mlflow.log_param("batch", batch)
    mlflow.log_param("epochs", epochs)

    # 加载模型
    model = YOLO("yolov8n.yaml")

    # 训练模型
    results = model.train(training_data=training_data, imgsz=imgsz, batch=batch, epochs=epochs)

    # 记录模型指标
    mlflow.log_metric("map", results.metrics["map"])

    # 保存模型文件
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=YOLOModel(),
        code_path=["predict_model.py"],
        artifacts={"weights": "path/to/weights.pt"},
        conda_env=None,
    )

if __name__ == "__main__":
    train()